import base64
import json
import os
from io import BytesIO

import cv2
import numpy as np
import torch
from PIL import Image
from flask import Flask, render_template, request, jsonify, send_from_directory

from model import CNN

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 2 * 1024 * 1024  # 2MB限制
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg'}

# 初始化模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)
model.load_state_dict(torch.load("best_model.pth", map_location=device))
model.eval()


def allowed_file(filename):
    """
    检查文件名是否包含扩展名（含.）
    提取扩展名后缀并转为小写
    判断小写扩展名是否在允许的扩展名集合ALLOWED_EXTENSIONS中
    """
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


# def preprocess_image(image):
#     """最基本的预处理图像适配MNIST输入"""
#     try:
#         img = image.convert('L').resize((28, 28))
#         img_array = np.array(img)
#         img_tensor = torch.from_numpy(img_array).float() / 255.0
#         img_tensor = img_tensor.unsqueeze(0).unsqueeze(0)  # [1,1,28,28]
#         return img_tensor.to(device)
#     except Exception as e:
#         raise ValueError(f"图像处理失败: {str(e)}")


def preprocess_image(image):
    """
    图像预处理函数（支持任意背景/数字颜色组合）

    :param image: 输入图像对象（PIL.Image格式）
    :returns: 预处理后的图像张量（Tensor格式，形状为[1,1,28,28]）
    :raises ValueError: 当图像处理过程失败时抛出异常

    处理流程：
    1. 转换为灰度图并扩展动态范围
    2. 执行动态颜色标准化
    3. 进行数字区域增强
    4. 尺寸标准化与对比度优化
    5. 转换为PyTorch张量并归一化
    """
    try:
        # 转换为灰度图并扩展动态范围
        img = image.convert('L')
        img_array = np.array(img)

        # 动态颜色标准化（自动处理明暗背景反转）
        img_array = dynamic_color_normalization(img_array)

        # 数字区域增强（定位并标准化数字主体区域）
        img_array = enhance_digit_region(img_array)

        # 尺寸标准化与对比度优化
        img_array = cv2.resize(img_array, (28, 28))
        img_array = contrast_enhancement(img_array)

        # 转换为张量并进行归一化处理
        img_tensor = torch.from_numpy(img_array).float() / 255.0
        img_tensor = img_tensor.unsqueeze(0).unsqueeze(0)
        return img_tensor.to(device)
    except Exception as e:
        raise ValueError(f"图像处理失败: {str(e)}")


# ------------ 核心子函数 ------------
def dynamic_color_normalization(img_array):
    """
    动态颜色通道优化（自动适配明暗背景）

    :param img_array: 输入图像数组（numpy.ndarray格式，形状为[H,W,C]或[H,W]）
    :returns: 颜色标准化后的图像数组（numpy.ndarray格式，形状为[H,W]）

    处理流程：
    1. 彩色图像转换到LAB颜色空间提取亮度通道
    2. 基于OTSU阈值判断背景类型
    3. 自动执行颜色反转保持黑底白字
    """
    # 自动选择最优颜色通道（针对彩色图像）
    if len(img_array.shape) == 3:  # 检查img_array的形状是否为三维，即是否为彩色图像（H,W,C）
        lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
        # 使用LAB空间的L通道（亮度）作为基准
        img_array = lab[:, :, 0]
    """
    使用cv2.threshold函数，参数是img_array，0阈值（因为OTSU会自动计算），255最大值，加上THRESH_BINARY和THRESH_OTSU标志。
    这会生成一个二值图像thresh。
    接着计算白色像素（255）和黑色像素（0）的数量，如果白色像素更多，说明背景是亮色的，需要将图像反转，用255减去原图像，这样就能得到黑底白字的效果。
    """
    # 动态反转判断（基于OTSU阈值）
    _, thresh = cv2.threshold(img_array, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    white_pixels = np.sum(thresh == 255)
    black_pixels = np.sum(thresh == 0)
    if white_pixels > black_pixels:
        img_array = 255 - img_array
    return img_array


def enhance_digit_region(img_array):
    """
    数字区域检测与增强（定位并标准化数字主体）

    :param img_array: 输入图像数组（numpy.ndarray格式，形状为[H,W]）
    :returns: 区域增强后的图像数组（numpy.ndarray格式，形状为[H,W]）

    处理流程：
    1. 多尺度边缘检测提取轮廓
    2. 形态学操作闭合孔洞
    3. 查找最大连通域确定数字区域
    4. 动态边界扩展与自适应缩放
    """
    # 多尺度边缘检测
    edges = cv2.Canny(img_array, 50, 150)

    # 形态学操作闭合小孔洞
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)

    # 查找最大连通域
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        max_contour = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(max_contour)

        # 动态扩展边界（基于原图尺寸比例）
        margin = int(max(w, h) * 0.2)
        x = max(0, x - margin)
        y = max(0, y - margin)
        w = min(img_array.shape[1] - x, w + 2 * margin)
        h = min(img_array.shape[0] - y, h + 2 * margin)
        roi = img_array[y:y + h, x:x + w]

        # 自适应缩放（保持纵横比）
        scaled = cv2.resize(roi, (20, 20))
        img_array = np.zeros((28, 28), dtype=np.uint8)
        img_array[4:24, 4:24] = scaled  # 居中放置
    return img_array


def contrast_enhancement(img_array):
    """
    多阶段对比度增强（提升数字笔画对比度）

    :param img_array: 输入图像数组（numpy.ndarray格式，形状为[H,W]）
    :returns: 对比度增强后的图像数组（numpy.ndarray格式，形状为[H,W]）

    处理流程：
    1. CLAHE局部对比度增强
    2. 全局直方图拉伸
    3. 非线性Gamma校正
    """
    # 阶段1：CLAHE局部对比度
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(4, 4))
    img_array = clahe.apply(img_array)

    # 阶段2：全局直方图拉伸
    p2, p98 = np.percentile(img_array, (2, 98))
    img_array = np.clip((img_array - p2) * 255.0 / (p98 - p2), 0, 255).astype(np.uint8)

    # 阶段3：非线性增强（Gamma校正）
    gamma = 1.5 if np.mean(img_array) < 127 else 0.7
    img_array = np.power(img_array / 255.0, gamma) * 255.0
    return img_array.astype(np.uint8)


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/static/<path:filename>')
def serve_static(filename):
    """
    使用send_from_directory安全读取static目录下的指定文件
    防止路径穿越攻击（如/static/../），确保只能访问静态资源目录内的文件
    """
    return send_from_directory('static', filename)


@app.route('/predict', methods=['POST'])
def predict():
    try:
        # 处理文件上传
        if 'file' in request.files:
            file = request.files['file']
            if not file or file.filename == '':
                return jsonify({'error': '未选择文件'}), 400
            if not allowed_file(file.filename):
                return jsonify({'error': '仅支持PNG/JPG/JPEG格式'}), 400

            image = Image.open(file.stream)
            tensor = preprocess_image(image)

        # 处理Canvas输入
        elif 'canvas' in request.form:
            img_data = request.form['canvas'].split(',')[1]  # 从请求中提取base64编码的图片数据，并去除前缀
            image = Image.open(BytesIO(base64.b64decode(img_data)))  # 解码base64数据为字节流，并转换为图像对象
            tensor = preprocess_image(image)

        else:
            return jsonify({'error': '未提供输入'}), 400

        with torch.no_grad():
            output = model(tensor.to(device))
            probabilities = torch.nn.functional.softmax(output[0], dim=0)
            pred = output.argmax(dim=1).item()
            confidence = probabilities[pred].item()

        return jsonify({
            'prediction': pred,
            'confidence': f"{confidence * 100:.1f}%"
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/logs')
def show_logs():
    if not os.path.exists("training_log.json"):
        return jsonify({'error': '未找到训练日志'}), 404

    try:
        with open("training_log.json", "r") as f:
            logs = []
            for line in f:
                try:
                    log_entry = json.loads(line.strip())
                    logs.append(log_entry)
                except json.JSONDecodeError:
                    continue  # 跳过无效行

        # 分离训练日志和测试结果
        train_logs = [log for log in logs if 'train_loss' in log]
        test_result = next((log for log in logs if log.get('type') == 'test_result'), None)

        if not train_logs:
            return jsonify({'error': '无有效训练日志'}), 404

        return jsonify({
            'best_val_loss': f"{min(log['val_loss'] for log in train_logs):.4f}",
            'final_lr': f"{train_logs[-1]['lr']:.2e}",
            'test_loss': f"{test_result['test_loss']:.4f}" if test_result else 'N/A',
            'test_acc': f"{test_result['test_acc']:.4%}" if test_result else 'N/A',
            'epochs': [log['epoch'] for log in train_logs],
            'train_loss': [log['train_loss'] for log in train_logs],
            'val_loss': [log['val_loss'] for log in train_logs],
            'train_acc': [log['train_acc'] for log in train_logs],
            'val_acc': [log['val_acc'] for log in train_logs],
            'confusion_matrix': '/static/confusion_matrix.png' if test_result else None
        })
    except Exception as e:
        return jsonify({'error': f'日志解析失败: {str(e)}'}), 500


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
